import cv2
import numpy as np
import signal
import sys
import requests
import json


def send_to_collector(url: str, stimulus):
    data = {
        'name': 'vsd',
        'stimulus': stimulus
    }

    try:

        response = requests.post(url, json=data)

        if response.status_code == 200 and json.loads(response.text)['code'] == 'OK':
            return True
        else:
            print(f'{response.status_code} and {response.text}')
            return False

    except requests.exceptions.ConnectionError as e:
        return False


# 设置计算平均识别率时考虑的最近检测次数
RECENT_DETECTION_COUNT = 5

# 加载类别名称
classNames = []
classFile = 'coco.names'
with open(classFile, 'rt') as f:
    classNames = f.read().rstrip('\n').split('\n')

# 加载模型
configPath = 'deploy.prototxt'
weightsPath ='mobilenet_iter_73000.caffemodel'

net = cv2.dnn_DetectionModel(weightsPath, configPath)
# 调整输入尺寸
net.setInputSize(224, 224)
net.setInputScale(1.0 / 127.5)
net.setInputMean((127.5, 127.5, 127.5))
net.setInputSwapRB(True)

# 尝试启用 GPU 加速
if cv2.cuda.getCudaEnabledDeviceCount() > 0:
    net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
    net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)


# 定义信号处理函数
def signal_handler(sig, frame):
    print('你按下了 Ctrl+C! 正在优雅退出...')
    cv2.destroyAllWindows()
    sys.exit(0)


# 注册信号处理函数
signal.signal(signal.SIGINT, signal_handler)


def visual_stimulus_detection(debug=False):
    # 打开摄像头
    cap = cv2.VideoCapture(0)

    frame_count = 0
    prev_objects = {}  # 存储上一帧物体信息，格式为 {classId: (box, disappear_count, appear_count, stable_count, seen, confidence_list)}
    seen_objects = set()  # 存储已经被“看见”过的物体的 classId

    # 原始的变化率阈值（基线）
    base_position_rate_threshold = 0.2
    base_size_rate_threshold = 0.2
    # 原始的物体连续消失次数阈值（基线）
    base_disappear_threshold = 6
    # 原始的物体连续出现次数阈值（基线）
    base_appear_threshold = base_disappear_threshold
    # 原始的物体稳定次数阈值（基线）
    base_stable_threshold = base_appear_threshold

    # 用于存储上一帧的刺激信息
    last_frame_stimulus = []

    while True:
        success, img = cap.read()
        if not success:
            break

        if frame_count % 2 == 0:
            # 降低置信度阈值
            classIds, confs, bbox = net.detect(img, confThreshold=0.1)

            current_objects = {}
            if len(classIds) != 0:
                for classId, confidence, box in zip(classIds.flatten(), confs.flatten(), bbox):
                    if classId in prev_objects:
                        prev_box, disappear_count, appear_count, stable_count, seen, confidence_list = prev_objects[classId]
                        appear_count += 1
                        disappear_count = 0
                        # 维护最近几次的置信度列表
                        confidence_list.append(confidence)
                        if len(confidence_list) > RECENT_DETECTION_COUNT:
                            confidence_list = confidence_list[-RECENT_DETECTION_COUNT:]

                        # 计算位置和大小的变化
                        center_prev = np.array([prev_box[0] + prev_box[2] // 2, prev_box[1] + prev_box[3] // 2])
                        center_current = np.array([box[0] + box[2] // 2, box[1] + box[3] // 2])
                        position_change = np.linalg.norm(center_prev - center_current)
                        size_change = abs(prev_box[2] * prev_box[3] - box[2] * box[3])

                        # 计算变化率
                        prev_size = prev_box[2] * prev_box[3]
                        size_change_rate = size_change / prev_size if prev_size != 0 else 0
                        prev_distance = np.linalg.norm(center_prev)
                        position_change_rate = position_change / prev_distance if prev_distance != 0 else 0

                        # 根据平均识别率动态调整变化率阈值
                        avg_confidence = np.mean(confidence_list)
                        position_rate_threshold = base_position_rate_threshold / avg_confidence
                        size_rate_threshold = base_size_rate_threshold / avg_confidence

                        if position_change_rate <= position_rate_threshold and size_change_rate <= size_rate_threshold:
                            stable_count += 1
                        else:
                            stable_count = 0

                    else:
                        appear_count = 1
                        disappear_count = 0
                        stable_count = 0
                        seen = False
                        confidence_list = [confidence]

                    # 计算该物体的平均识别率
                    avg_confidence = np.mean(confidence_list)

                    # 根据平均识别率动态调整变化率阈值
                    position_rate_threshold = base_position_rate_threshold / avg_confidence
                    size_rate_threshold = base_size_rate_threshold / avg_confidence

                    # 根据平均识别率动态调整看见、消失和稳定的阈值
                    appear_threshold = int(base_appear_threshold / avg_confidence)
                    disappear_threshold = int(base_disappear_threshold / avg_confidence)
                    stable_threshold = int(base_stable_threshold / avg_confidence)

                    cv2.rectangle(img, box, color=(0, 255, 0), thickness=2)
                    # 可以根据需要决定是否显示文本，减少不必要的操作
                    cv2.putText(img, classNames[classId - 1].upper(), (box[0] + 10, box[1] + 30),
                                cv2.FONT_HERSHEY_COMPLEX, 1, (0, 255, 0), 2)
                    cv2.putText(img, str(round(confidence * 100, 2)) + '%', (box[0] + 200, box[1] + 30),
                                cv2.FONT_HERSHEY_COMPLEX, 1, (0, 255, 0), 2)

                    if stable_count == stable_threshold and not seen:
                        center = (int(box[0] + box[2] // 2), int(box[1] + box[3] // 2))
                        last_frame_stimulus.append(f"刺激，物体 {classNames[classId - 1]} 首次被看见了，位置: {center}")
                        seen = True
                        seen_objects.add(classId)

                    current_objects[classId] = (box, disappear_count, appear_count, stable_count, seen, confidence_list)

            # 检查物体变化
            for prev_classId, (prev_box, disappear_count, appear_count, stable_count, seen, confidence_list) in prev_objects.items():
                if prev_classId in current_objects:
                    current_box = current_objects[prev_classId][0]
                    # 计算该物体的平均识别率
                    avg_confidence = np.mean(confidence_list)

                    # 根据平均识别率动态调整变化率阈值
                    position_rate_threshold = base_position_rate_threshold / avg_confidence
                    size_rate_threshold = base_size_rate_threshold / avg_confidence

                    # 根据平均识别率动态调整看见、消失和稳定的阈值
                    appear_threshold = int(base_appear_threshold / avg_confidence)
                    disappear_threshold = int(base_disappear_threshold / avg_confidence)
                    stable_threshold = int(base_stable_threshold / avg_confidence)

                    # 计算位置和大小的变化
                    center_prev = np.array([prev_box[0] + prev_box[2] // 2, prev_box[1] + prev_box[3] // 2])
                    center_current = np.array([current_box[0] + current_box[2] // 2, current_box[1] + current_box[3] // 2])
                    position_change = np.linalg.norm(center_prev - center_current)
                    size_change = abs(prev_box[2] * prev_box[3] - current_box[2] * current_box[3])

                    # 计算变化率
                    prev_size = prev_box[2] * prev_box[3]
                    size_change_rate = size_change / prev_size if prev_size != 0 else 0
                    prev_distance = np.linalg.norm(center_prev)
                    position_change_rate = position_change / prev_distance if prev_distance != 0 else 0

                    # 仅当位置或大小变化率超过阈值且物体连续出现次数和稳定次数达到阈值时才记录信息
                    if (position_change_rate > position_rate_threshold or size_change_rate > size_rate_threshold) and appear_count >= appear_threshold and stable_count >= stable_threshold:
                        center_current_int = (int(center_current[0]), int(center_current[1]))
                        last_frame_stimulus.append(f"受到刺激，物体 {classNames[prev_classId - 1]} 位置变化率: {position_change_rate:.2f}，大小变化率: {size_change_rate:.2f}，当前位置: {center_current_int}")
                else:
                    # 计算该物体的平均识别率
                    avg_confidence = np.mean(confidence_list)

                    # 根据平均识别率动态调整消失的阈值
                    disappear_threshold = int(base_disappear_threshold / avg_confidence)

                    # 物体未在当前帧出现，增加消失次数，重置出现次数和稳定次数
                    disappear_count += 1
                    appear_count = 0
                    stable_count = 0

                    # 只有当物体已经被“看见”过，且连续消失次数超过阈值时，才触发消失的刺激
                    if seen and disappear_count > disappear_threshold:
                        last_position = (int(prev_box[0] + prev_box[2] // 2), int(prev_box[1] + prev_box[3] // 2))
                        last_frame_stimulus.append(f"受到刺激，物体 {classNames[prev_classId - 1]} 连续消失超过 {disappear_threshold} 次，最后位置: {last_position}")
                    else:
                        current_objects[prev_classId] = (prev_box, disappear_count, appear_count, stable_count, seen, confidence_list)

            prev_objects = current_objects
        else:
            for stimulus in last_frame_stimulus:
                # 打印上一帧的刺激信息
                print(stimulus)
                # 发送给收集刺激的服务
                send_to_collector('http://localhost:8002/regist', stimulus)

            last_frame_stimulus = []

        frame_count += 1
        if debug:
            cv2.imshow("Output", img)
            if cv2.waitKey(1) == ord('q'):
                break

    cap.release()
    if debug:
        cv2.destroyAllWindows()


if __name__ == "__main__":
    visual_stimulus_detection(debug=False)
    